#!/usr/bin/env python
# Author: mingzhong.huang

import requests
from requests.packages.urllib3.exceptions import InsecureRequestWarning
import json
from distutils.util import strtobool
from setting import config
from core.logger import logger
from sqlalchemy.orm import Session
from pprint import pprint





class K8sBlackBoxProbe(object):
    def __init__(self):
        self.url = config.K8S_API
        self.token = config.K8S_TOKEN
        self.namespace = config.K8S_PROM_NAMESPACE
        self.header = {
           "Authorization":"Bearer "+self.token,
           "Content-Type": "application/json"
        }
        requests.packages.urllib3.disable_warnings(InsecureRequestWarning)

    def get_blackbox_all_probes(self):
        """
        获取所有的probe探针
        """
        try:
            url = self.url + f"/apis/monitoring.coreos.com/v1/namespaces/{self.namespace}/probes"
            req = requests.get(url=url, headers=self.header, verify=False)
            if req.status_code != 200:
                logger.error(f"get blackbox all probes failed => {req.text}")
                return False
            blackbox_probes = json.loads(req.text)
        except Exception as e:
            logger.error(f"get blackbox all probes error => {e}")
            return False

        logger.info(f"get blackbox all probes success")
        return blackbox_probes["items"]


    def get_blackbox_probe(self, name):
        """
        获取一个probe探针
        """
        try:
            url = self.url + f"/apis/monitoring.coreos.com/v1/namespaces/{self.namespace}/probes/{name}"
            req = requests.get(url=url, headers=self.header, verify=False)
            if req.status_code != 200:
                logger.error(f"get blackbox probe:{name} failed => {req.text}")
                return False
            blackbox_probe = json.loads(req.text)
        except Exception as e:
            logger.error(f"get blackbox probe:{name} error => {e}")
            return False

        logger.info(f"get blackbox probe:{name} success")
        return blackbox_probe


    def is_exist(self, name):
        """
        获取一个probe探针
        """
        try:
            url = self.url + f"/apis/monitoring.coreos.com/v1/namespaces/{self.namespace}/probes/{name}"
            req = requests.get(url=url, headers=self.header, verify=False)
            if req.status_code != 200:
                return False
        except Exception as e:
            logger.error(f"is_exist blackbox probe:{name} error => {e}")
            return False

        return True



    def add_blackbox_probe(self, probe_info):
        """
        增加一个probe探针
        """
        try:
            # 判断probe探针是否存在
            if self.is_exist(probe_info['name']):
                logger.error(f"add blackbox probe:{probe_info['name']} failed, was exist")
                return False

            probe_data = self.convert_probe_data(probe_info)

            data = json.dumps(probe_data, ensure_ascii=False).encode()
            url = self.url + f"/apis/monitoring.coreos.com/v1/namespaces/{self.namespace}/probes"
            req = requests.post(url=url, headers=self.header, verify=False, data=data)
            if req.status_code not in [200, 201]:
                logger.error(f"add blackbox probe:{probe_info['name']} failed => {req.text}")
                return False
        except Exception as e:
            logger.error(f"add blackbox probe:{probe_info['name']} error => {e}")
            return False

        logger.info(f"add blackbox probe:{probe_info['name']} success")
        return True



    def update_blackbox_probe(self, probe_info):
        """
        更新一个probe探针
        """
        try:
            # 拿到现在k8s集群中的probe资源的版本号
            k8s_blackbox_probe = self.get_blackbox_probe(probe_info['name'])
            if not k8s_blackbox_probe:
                logger.error(f"get blackbox probe:{probe_info['name']} failed, not exist")
                return False

            # 更新put请求的数据
            probe_data = self.convert_probe_data(probe_info)
            probe_data["metadata"]["resourceVersion"] = k8s_blackbox_probe["metadata"]["resourceVersion"]
            data = json.dumps(probe_data, ensure_ascii=False).encode()

            # 更新probe
            url = self.url + f"/apis/monitoring.coreos.com/v1/namespaces/{self.namespace}/probes/{probe_info['name']}"
            req = requests.put(url=url, data=data, headers=self.header, verify=False)
            if req.status_code != 200:
                logger.error(f"update blackbox probe:{probe_info['name']} failed => {req.text}")
                return False

        except Exception as e:
            logger.error(f"update blackbox probe:{probe_info['name']} error => {e}")
            return False

        logger.info(f"update blackbox probe:{probe_info['name']} success")
        return True


    def delete_blackbox_probe(self, name):
        """
        删除一个probe探针
        """
        try:
            url = self.url + f"/apis/monitoring.coreos.com/v1/namespaces/{self.namespace}/probes/{name}"
            req = requests.delete(url=url, headers=self.header, verify=False)
            if req.status_code != 200:
                logger.error(f"delete blackbox probe:{name} failed => {req.text}")
                return False
        except Exception as e:
            logger.error(f"delete blackbox probe:{name} error => {e}")
            return False

        logger.info(f"delete blackbox probe:{name} success")
        return True



    def convert_probe_data(self, probe_info: dict):
        module_data = {
            "http": "http_2xx",
            "tcp": "tcp_connect",
            "icmp": "icmp"
        }
        probe_data={
            "apiVersion": 'monitoring.coreos.com/v1',
            "kind": 'Probe',
            "metadata":{
                "name": probe_info["name"],
                "namespace": self.namespace,
            },
            "spec": {
                "interval": '60s',
                "jobName": probe_info["name"],
                "module": module_data[probe_info["module"]],
                'prober': {
                    'path': '/probe',
                    'scheme': 'http',
                    'url': 'blackbox-exporter.monitoring.svc:19115'
                },
                'scrapeTimeout': '30s',
                'targets': {
                    'staticConfig': {
                        'labels': probe_info["labels"],
                        'static': probe_info["targets"]
                    }
                }
            }
        }

        return probe_data


def main():
        blackbox_probe = K8sBlackBoxProbe()
        # blackbox_probes = blackbox_probe.get_blackbox_all_probes()
        # blackbox_probe = blackbox_probe.get_blackbox_probe("yunwei-k8s-master2")
        # pprint(blackbox_probe)

        tcp_probe_info = {
            "name": "yunwei-k8s-master3",
            "module": "tcp_connect",
            "labels": {
                'i_blackbox_group': 'yunwei-k8s',
                'i_blackbox_name': 'yunwei-k8s-master',
                'i_department': '运维'
            },
            "targets": [
                '10.4.0.21:10251',
                '10.4.0.21:22',
            ]
        }
        # result = blackbox_probe.add_blackbox_probe(tcp_probe_info)
        # result = blackbox_probe.update_blackbox_probe(tcp_probe_info)
        # print(result)
        # blackbox_probe.delete_blackbox_probe(tcp_probe_info["name"])

        http_probe_info = {
            "name": "yunwei-url-test",
            "module": "http_2xx",
            "labels": {
                'i_blackbox_group': 'yunwei-url',
                'i_blackbox_name': 'yunwei-url-test',
                'i_department': '运维'
            },
            "targets": [
                'https://www.baidu.com/abcd'
            ]
        }
        # result = blackbox_probe.add_blackbox_probe(http_probe_info)
        # result = blackbox_probe.update_blackbox_probe(http_probe_info)
        # print(result)
        blackbox_probe.delete_blackbox_probe(http_probe_info["name"])

if __name__ == "__main__":
    main()

